#!/usr/bin/env python3

import argparse
import json
import rich
from rich.console import Console
from rich.syntax import Syntax
from rich.tree import Tree

class AgentParser:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description="A command line parser for Chatbot Agent")
        self._setup_arguments()

    def _setup_arguments(self):

        self.parser.add_argument(
            "--model",
            type=str,
            help="Path of HuggingFace model weights, e.g., '/home/wangxianda/model_zoo/llama_3.2_3b_instruct'",
            required=True
        )
        self.parser.add_argument(
            "--device",
            type=str,
            help="Device to run the model on. Use 'cuda' for GPU or 'cpu' for CPU. Default is 'cuda'.",
            default="cuda"
        )
        self.parser.add_argument(
            "--dtype",
            type=str,
            help="Data type for model weights. Use 'torch.float16' for half-precision or 'torch.float32' for full-precision. Default is 'torch.float16'.",
            default="float16"
        )
        # self.parser.add_argument(
        #     "--context_len",
        #     type=int,
        #     help="History context length included in one chat. Default is '10'.",
        #     default="10"
        # )
        # self.parser.add_argument(
        #     "--max_tokens",
        #     type=int,
        #     help="Number of max tokens included in one response. Default is '200'.",
        #     default="200"
        # )

    def parse(self):
        return self.parser.parse_args()

    def spec(self):
        # Parse command-line arguments
        args = self.parser.parse_args()

        # Define Monokai colors
        tree_style = "bold white"     # For the first level (root)
        branch_style = "bold cyan"     # For the second level
        leaf_style = "dim magenta"     # For the leaf nodes (values)

        # Create the tree with a Monokai-like background (dark)
        console = Console()  # Set the console background to black
        tree = Tree("Parsed Arguments", style=tree_style, guide_style=tree_style)

        # Iterate over parsed arguments and add them to the tree
        for key, value in vars(args).items():
            branch = tree.add(key, style=branch_style, guide_style=branch_style)

            # Format the value to display
            if isinstance(value, dict):
                # If the value is a dictionary, convert it to a string for display
                branch_content = json.dumps(value, indent=4)
                branch.add(Syntax(branch_content, "json", theme="monokai"))
            else:
                # Otherwise, just add the value directly
                branch.add(str(value), style=leaf_style)

        # Print the parsing message and tree structure
        rich.print("\n:raccoon: Parsing Arguments ... \n")
        console.print(tree)

if __name__ == "__main__":
    parser = AgentParser()
    args = parser.parse()
    print(f"{args}")
